import os
import sys
from frb_pipeline_tools import *

if len(sys.argv) < 4:
    print("❌ 用法: python run_frb_search.py /path/to/fits /path/to/filtered.txt /path/to/plots [--save_npz]")
    sys.exit(1)

fits_dir = sys.argv[1]
cand_txt_path = sys.argv[2]
plot_output_dir = sys.argv[3]
save_npz = "--save_npz" in sys.argv

npz_dir = os.path.join(os.path.dirname(plot_output_dir), "npz")
os.makedirs(plot_output_dir, exist_ok=True)
if save_npz:
    os.makedirs(npz_dir, exist_ok=True)

fits_list = sorted([
    os.path.join(fits_dir, f) for f in os.listdir(fits_dir)
    if f.endswith(".fits")
])
fits_index_map = {os.path.basename(f): i for i, f in enumerate(fits_list)}

with open(cand_txt_path, 'r') as f:
    lines = [line.strip() for line in f if line.strip()]
print(f"📋 共有 {len(lines)} 个候选体")

for i, line in enumerate(lines):
    try:
        parts = line.split()
        fits_name = parts[0]
        sample_id = int(parts[1])
        DM = float(parts[2])
        snr = float(parts[-1])

        if fits_name not in fits_index_map:
            print(f"⚠️ 跳过：未找到 {fits_name}")
            continue

        index = fits_index_map[fits_name]
        fits_path = fits_list[index]

        data, freqs, times, tsamp, mjd_start = extract_candidate_segment(
            fits_list, index, sample_id, DM)
        data_clean = normalise(data)
        data_dedisp = dedisperse(data_clean, freqs, tsamp, DM)

        base = os.path.splitext(fits_name)[0]
        out_png = os.path.join(plot_output_dir, f"{base}_{sample_id:06d}_DM{int(DM)}.png")


        plot_frb_four_panel_summary(
            data_clean, data_dedisp, freqs, times, tsamp,
            snr, DM, sample_id, fits_path, mjd_start, i+1, outname=out_png
        )

        if save_npz:
            npz_path = os.path.join(npz_dir, base + ".npz")
            np.savez(npz_path, data=data, freqs=freqs, times=times,
                     data_dedisp=data_dedisp)

    except Exception as e:
        print(f"[❌] 第 {i+1} 行出错: {e}")
        continue

print("✅ 所有图像已输出至", plot_output_dir)